
#include <iostream>
#include <set>
#include <map>

using namespace std;

#include "Mutex.h"
#include "Cond.h"
#include "Lockable.h"

Lockable::Lockable () {
 readLockedCond.setMutex  (readLockedMutex);
 writeLockedCond.setMutex (writeLockedMutex);
 writerCountCond.setMutex (writerCountMutex);
 readerCountCond.setMutex (readerCountMutex);
 readLocked = false;
 writeLocked = false;
 writerCount = 0;
 readerCount = 0;
}

bool Lockable::readLock () {
 readLockedMutex.lock ();
 if ( readLocked ) { /* it is already locked, wait until it is not */
  readLockedCond.wait ();
 }
 readLocked = true; /* No new threads will attach as readers or writers after this */
 readLockedMutex.unlock ();
  
 writerCountMutex.lock ();
 while (writerCount > 0) { /* wait until all writers finish their jobs */
  writerCountCond.wait ();
 }
 writerCountMutex.unlock ();

 readerCountMutex.lock ();  
 while (readerCount > 0) { /* wait until all readers finish their jobs */
  readerCountCond.wait ();
 }
 readerCountMutex.unlock ();

 return true;
}

bool Lockable::writeLock () {
 
 writeLockedMutex.lock ();
 if ( writeLocked ) { /* already locked */
  writeLockedCond.wait();
 }
 writeLocked = true;
 writeLockedMutex.unlock ();

 writerCountMutex.lock ();
 while (writerCount > 0) { /* wait until all writers finish their jobs */
  writerCountCond.wait ();
 }
 writerCount = 1; /* This thread is the only writer */
 writerCountMutex.unlock ();

 return true;
}

bool Lockable::writeLock (int rowNum, bool block) {
 writeLockedMutex.lock ();

 if ( writeLocked ) {
  writeLockedCond.wait ();
 }

 writeLockedRowsMutex.lock ();
 set<int>::iterator writeLockItr = writeLockedRows.find (rowNum);

 if ( writeLockItr != writeLockedRows.end () ) { /* row is already locked */
  if (block) { /* blocking mode, wait for it to be unlocked */
   writeLockedRowsCond[ rowNum ].wait ();
  }
  else {
   return false;
  }
 }

 writerCountMutex.lock ();
 writerCount++; /* no one can writeLock after this */
 writerCountMutex.unlock ();
 writeLockedMutex.unlock ();

 Cond cond;
 cond.setMutex ( writeLockedRowsMutex );
 
 writeLockedRows.insert ( rowNum );
 writeLockedRowsCond.insert ( pair<int, Cond> ( rowNum, cond ) ); 

 writeLockedRowsMutex.unlock ();
 return true;
}

bool Lockable::readUnlock () {
 readLockedMutex.lock ();
 readLocked = false;
 readLockedCond.broadcast ();
 readLockedMutex.unlock ();
 return true;
}

bool Lockable::writeUnlock () {
 writeLockedMutex.lock ();
 writeLocked = false;
 
 writerCountMutex.lock ();
 writerCount = 0;
 writerCountCond.broadcast ();
 writerCountMutex.unlock ();
 
 writeLockedCond.broadcast ();
 writeLockedMutex.unlock ();
 return true;
}

bool Lockable::writeUnlock (int rowNum) {
 writeLockedRowsMutex.lock ();
 set<int>::iterator writeLockItr = writeLockedRows.find (rowNum);

 if ( writeLockItr == writeLockedRows.end () ) { /* row is not locked! */
  writeLockedRowsMutex.unlock ();
  return false;
 }
 else { /* row is locked */
  writeLockedRows.erase ( writeLockItr );
  writeLockedRowsCond[ rowNum ].broadcast();
  writeLockedRowsCond.erase ( rowNum );
  writeLockedRowsMutex.unlock ();

  writerCountMutex.lock ();
  writerCount--;
  writerCountCond.broadcast ();
  writerCountMutex.unlock ();

  return true;
 }
}

void Lockable::attachReader () { 
 readLockedMutex.lock ();
 if ( readLocked ) {
  readLockedCond.wait ();
 }
 readerCountMutex.lock ();
 readerCount++;
 readerCountMutex.unlock ();
 readLockedMutex.unlock ();
}

void Lockable::detachReader () {
 readerCountMutex.lock ();
 readerCount--;
 readerCountCond.broadcast ();
 readerCountMutex.unlock ();
}

